# coding=utf-8

import base64
import json
from Crypto.Cipher import AES
from binascii import b2a_hex, a2b_hex


class WXBizDataCrypt:
    def __init__(self, appId, sessionKey):
        self.appId = appId
        self.sessionKey = sessionKey

    def decrypt(self, encryptedData, iv):
        # base64 decode
        sessionKey = base64.b64decode(self.sessionKey)
        encryptedData = base64.b64decode(encryptedData)
        iv = base64.b64decode(iv)

        cipher = AES.new(sessionKey, AES.MODE_CBC, iv)

        decrypted = json.loads(self._unpad(cipher.decrypt(encryptedData)))

        if decrypted['watermark']['appid'] != self.appId:
            raise Exception('Invalid Buffer')

        return decrypted

    def _unpad(self, s):
        return s[:-ord(s[len(s)-1:])]

    # def decrypt1(self, encryptedData, iv):
    #     # base64 decode
    #     sessionKey = self.sessionKey.encode('utf-8')
    #     encryptedData = base64.b64decode(encryptedData.encode('utf8'))
    #     iv = iv.encode('utf-8')
    #     cipher = AES.new(sessionKey, AES.MODE_CBC, iv)
    #     decrypted = json.loads(bytes.decode(cipher.decrypt(encryptedData)).rstrip('\0'))
    #
    #     if decrypted['watermark']['appid'] != self.appId:
    #         raise Exception('Invalid Buffer')
    #
    #     return decrypted
    #
    # def add_to_16(self, text):
    #     if len(text.encode('utf-8')) % 16:
    #         add = 16 - (len(text.encode('utf-8')) % 16)
    #     else:
    #         add = 0
    #     text = text + ('\0' * add)
    #     return text.encode('utf-8')
